{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "start_time": "2023-07-08T10:54:43.346709Z",
     "end_time": "2023-07-08T10:54:48.972209Z"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import gym\n",
    "import numpy as np\n",
    "import collections\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "import rl_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "class VAnet(torch.nn.Module):\n",
    "    ''' 只有一层隐藏层的A网络和V网络 '''\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim):\n",
    "        super(VAnet, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)  # 共享网络部分\n",
    "        self.fc_A = torch.nn.Linear(hidden_dim, action_dim)\n",
    "        self.fc_V = torch.nn.Linear(hidden_dim, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        A = self.fc_A(F.relu(self.fc1(x)))\n",
    "        V = self.fc_V(F.relu(self.fc1(x)))\n",
    "        Q = V + A - A.mean(1).view(-1, 1)  # Q值由V值和A值计算得到\n",
    "        return Q"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2023-07-08T10:54:48.975165Z",
     "end_time": "2023-07-08T10:54:48.987101Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "class Qnet(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    只有一层隐藏层的Q网络\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim):\n",
    "        super().__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        return self.fc2(x)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2023-07-08T10:54:48.988099Z",
     "end_time": "2023-07-08T10:54:49.006082Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "class DQN:\n",
    "    ''' DQN算法,包括Double DQN和Dueling DQN '''\n",
    "    def __init__(self,\n",
    "                 state_dim,\n",
    "                 hidden_dim,\n",
    "                 action_dim,\n",
    "                 learning_rate,\n",
    "                 gamma,\n",
    "                 epsilon,\n",
    "                 target_update,\n",
    "                 device,\n",
    "                 dqn_type='VanillaDQN'):\n",
    "        self.action_dim = action_dim\n",
    "        if dqn_type == 'DuelingDQN':  # Dueling DQN采取不一样的网络框架\n",
    "            self.q_net = VAnet(state_dim, hidden_dim,\n",
    "                               self.action_dim).to(device)\n",
    "            self.target_q_net = VAnet(state_dim, hidden_dim,\n",
    "                                      self.action_dim).to(device)\n",
    "        else:\n",
    "            self.q_net = Qnet(state_dim, hidden_dim,\n",
    "                              self.action_dim).to(device)\n",
    "            self.target_q_net = Qnet(state_dim, hidden_dim,\n",
    "                                     self.action_dim).to(device)\n",
    "        self.optimizer = torch.optim.Adam(self.q_net.parameters(),\n",
    "                                          lr=learning_rate)\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = epsilon\n",
    "        self.target_update = target_update\n",
    "        self.count = 0\n",
    "        self.dqn_type = dqn_type\n",
    "        self.device = device\n",
    "\n",
    "    def take_action(self, state):\n",
    "        if np.random.random() < self.epsilon:\n",
    "            action = np.random.randint(self.action_dim)\n",
    "        else:\n",
    "            state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
    "            action = self.q_net(state).argmax().item()\n",
    "        return action\n",
    "\n",
    "    def max_q_value(self, state):\n",
    "        state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
    "        return self.q_net(state).max().item()\n",
    "\n",
    "    def update(self, transition_dict):\n",
    "        states = torch.tensor(transition_dict['states'],\n",
    "                              dtype=torch.float).to(self.device)\n",
    "        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(\n",
    "            self.device)\n",
    "        rewards = torch.tensor(transition_dict['rewards'],\n",
    "                               dtype=torch.float).view(-1, 1).to(self.device)\n",
    "        next_states = torch.tensor(transition_dict['next_states'],\n",
    "                                   dtype=torch.float).to(self.device)\n",
    "        dones = torch.tensor(transition_dict['dones'],\n",
    "                             dtype=torch.float).view(-1, 1).to(self.device)\n",
    "\n",
    "        q_values = self.q_net(states).gather(1, actions)\n",
    "        if self.dqn_type == 'DoubleDQN':\n",
    "            max_action = self.q_net(next_states).max(1)[1].view(-1, 1)\n",
    "            max_next_q_values = self.target_q_net(next_states).gather(\n",
    "                1, max_action)\n",
    "        else:\n",
    "            max_next_q_values = self.target_q_net(next_states).max(1)[0].view(\n",
    "                -1, 1)\n",
    "        q_targets = rewards + self.gamma * max_next_q_values * (1 - dones)\n",
    "        dqn_loss = torch.mean(F.mse_loss(q_values, q_targets))\n",
    "        self.optimizer.zero_grad()\n",
    "        dqn_loss.backward()\n",
    "        self.optimizer.step()\n",
    "\n",
    "        if self.count % self.target_update == 0:\n",
    "            self.target_q_net.load_state_dict(self.q_net.state_dict())\n",
    "        self.count += 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2023-07-08T10:54:49.003091Z",
     "end_time": "2023-07-08T10:54:49.037002Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "def dis_to_con(discrete_action, env, action_dim):\n",
    "    \"\"\"\n",
    "    离散动作转回连续的函数\n",
    "    :param discrete_action:\n",
    "    :param env:\n",
    "    :param action_dim:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    # 连续动作的最小值\n",
    "    action_lowbound = env.action_space.low[0]\n",
    "    # 连续动作的最大值\n",
    "    action_upbound = env.action_space.high[0]\n",
    "    return action_lowbound + (discrete_action / (action_dim - 1)) * (action_upbound - action_lowbound)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2023-07-08T10:54:49.036003Z",
     "end_time": "2023-07-08T10:54:49.049984Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "def train_DQN(agent, env, num_episodes, replay_buffer, minimal_size,\n",
    "              batch_size):\n",
    "    return_list = []\n",
    "    max_q_value_list = []\n",
    "    max_q_value = 0\n",
    "    for i in range(10):\n",
    "        with tqdm(total=int(num_episodes / 10),\n",
    "                  desc='Iteration %d' % i) as pbar:\n",
    "            for i_episode in range(int(num_episodes / 10)):\n",
    "                episode_return = 0\n",
    "                state = env.reset()\n",
    "                done = False\n",
    "                while not done:\n",
    "                    action = agent.take_action(state)\n",
    "                    max_q_value = agent.max_q_value(\n",
    "                        state) * 0.005 + max_q_value * 0.995  # 平滑处理\n",
    "                    max_q_value_list.append(max_q_value)  # 保存每个状态的最大Q值\n",
    "                    # Double DQN的区别\n",
    "                    action_continuous = dis_to_con(action, env,\n",
    "                                                   agent.action_dim)\n",
    "                    next_state, reward, done, _ = env.step([action_continuous])\n",
    "                    replay_buffer.add(state, action, reward, next_state, done)\n",
    "                    state = next_state\n",
    "                    episode_return += reward\n",
    "                    if replay_buffer.size() > minimal_size:\n",
    "                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(\n",
    "                            batch_size)\n",
    "                        transition_dict = {\n",
    "                            'states': b_s,\n",
    "                            'actions': b_a,\n",
    "                            'next_states': b_ns,\n",
    "                            'rewards': b_r,\n",
    "                            'dones': b_d\n",
    "                        }\n",
    "                        agent.update(transition_dict)\n",
    "                return_list.append(episode_return)\n",
    "                if (i_episode + 1) % 10 == 0:\n",
    "                    pbar.set_postfix({\n",
    "                        'episode':\n",
    "                        '%d' % (num_episodes / 10 * i + i_episode + 1),\n",
    "                        'return':\n",
    "                        '%.3f' % np.mean(return_list[-10:])\n",
    "                    })\n",
    "                pbar.update(1)\n",
    "    return return_list, max_q_value_list"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2023-07-08T10:54:49.051960Z",
     "end_time": "2023-07-08T10:54:49.071909Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "E:\\PCL\\PyCharm 2023.1\\jbr\\bin\\D\\code\\env\\lib\\site-packages\\gym\\envs\\registration.py:593: UserWarning: \u001B[33mWARN: The environment CartPole-v0 is out of date. You should consider upgrading to version `v1`.\u001B[0m\n",
      "  logger.warn(\n",
      "E:\\PCL\\PyCharm 2023.1\\jbr\\bin\\D\\code\\env\\lib\\site-packages\\gym\\core.py:317: DeprecationWarning: \u001B[33mWARN: Initializing wrapper in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001B[0m\n",
      "  deprecation(\n",
      "E:\\PCL\\PyCharm 2023.1\\jbr\\bin\\D\\code\\env\\lib\\site-packages\\gym\\wrappers\\step_api_compatibility.py:39: DeprecationWarning: \u001B[33mWARN: Initializing environment in old step API which returns one bool instead of two. It is recommended to set `new_step_api=True` to use new step API. This will be the default behaviour in future.\u001B[0m\n",
      "  deprecation(\n",
      "E:\\PCL\\PyCharm 2023.1\\jbr\\bin\\D\\code\\env\\lib\\site-packages\\gym\\core.py:256: DeprecationWarning: \u001B[33mWARN: Function `env.seed(seed)` is marked as deprecated and will be removed in the future. Please use `env.reset(seed=seed)` instead.\u001B[0m\n",
      "  deprecation(\n",
      "Iteration 0:   0%|          | 0/50 [00:00<?, ?it/s]C:\\Users\\14227\\AppData\\Local\\Temp\\ipykernel_9536\\2837068133.py:37: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ..\\torch\\csrc\\utils\\tensor_new.cpp:248.)\n",
      "  state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
      "Iteration 0:   0%|          | 0/50 [00:05<?, ?it/s]\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'Discrete' object has no attribute 'low'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mAttributeError\u001B[0m                            Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[7], line 22\u001B[0m\n\u001B[0;32m     19\u001B[0m replay_buffer \u001B[38;5;241m=\u001B[39m rl_utils\u001B[38;5;241m.\u001B[39mReplayBuffer(buffer_size)\n\u001B[0;32m     20\u001B[0m agent \u001B[38;5;241m=\u001B[39m DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon,\n\u001B[0;32m     21\u001B[0m             target_update, device, \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mDuelingDQN\u001B[39m\u001B[38;5;124m'\u001B[39m)\n\u001B[1;32m---> 22\u001B[0m return_list, max_q_value_list \u001B[38;5;241m=\u001B[39m \u001B[43mtrain_DQN\u001B[49m\u001B[43m(\u001B[49m\u001B[43magent\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43menv\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mnum_episodes\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m     23\u001B[0m \u001B[43m                                          \u001B[49m\u001B[43mreplay_buffer\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mminimal_size\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m     24\u001B[0m \u001B[43m                                          \u001B[49m\u001B[43mbatch_size\u001B[49m\u001B[43m)\u001B[49m\n",
      "Cell \u001B[1;32mIn[6], line 19\u001B[0m, in \u001B[0;36mtrain_DQN\u001B[1;34m(agent, env, num_episodes, replay_buffer, minimal_size, batch_size)\u001B[0m\n\u001B[0;32m     17\u001B[0m max_q_value_list\u001B[38;5;241m.\u001B[39mappend(max_q_value)  \u001B[38;5;66;03m# 保存每个状态的最大Q值\u001B[39;00m\n\u001B[0;32m     18\u001B[0m \u001B[38;5;66;03m# Double DQN的区别\u001B[39;00m\n\u001B[1;32m---> 19\u001B[0m action_continuous \u001B[38;5;241m=\u001B[39m \u001B[43mdis_to_con\u001B[49m\u001B[43m(\u001B[49m\u001B[43maction\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43menv\u001B[49m\u001B[43m,\u001B[49m\n\u001B[0;32m     20\u001B[0m \u001B[43m                               \u001B[49m\u001B[43magent\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43maction_dim\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m     21\u001B[0m next_state, reward, done, _ \u001B[38;5;241m=\u001B[39m env\u001B[38;5;241m.\u001B[39mstep([action_continuous])\n\u001B[0;32m     22\u001B[0m replay_buffer\u001B[38;5;241m.\u001B[39madd(state, action, reward, next_state, done)\n",
      "Cell \u001B[1;32mIn[5], line 10\u001B[0m, in \u001B[0;36mdis_to_con\u001B[1;34m(discrete_action, env, action_dim)\u001B[0m\n\u001B[0;32m      2\u001B[0m \u001B[38;5;250m\u001B[39m\u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[0;32m      3\u001B[0m \u001B[38;5;124;03m离散动作转回连续的函数\u001B[39;00m\n\u001B[0;32m      4\u001B[0m \u001B[38;5;124;03m:param discrete_action:\u001B[39;00m\n\u001B[1;32m   (...)\u001B[0m\n\u001B[0;32m      7\u001B[0m \u001B[38;5;124;03m:return:\u001B[39;00m\n\u001B[0;32m      8\u001B[0m \u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[0;32m      9\u001B[0m \u001B[38;5;66;03m# 连续动作的最小值\u001B[39;00m\n\u001B[1;32m---> 10\u001B[0m action_lowbound \u001B[38;5;241m=\u001B[39m \u001B[43menv\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43maction_space\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mlow\u001B[49m[\u001B[38;5;241m0\u001B[39m]\n\u001B[0;32m     11\u001B[0m \u001B[38;5;66;03m# 连续动作的最大值\u001B[39;00m\n\u001B[0;32m     12\u001B[0m action_upbound \u001B[38;5;241m=\u001B[39m env\u001B[38;5;241m.\u001B[39maction_space\u001B[38;5;241m.\u001B[39mhigh[\u001B[38;5;241m0\u001B[39m]\n",
      "\u001B[1;31mAttributeError\u001B[0m: 'Discrete' object has no attribute 'low'"
     ]
    }
   ],
   "source": [
    "lr = 2e-3\n",
    "num_episodes = 500\n",
    "hidden_dim = 128\n",
    "gamma = 0.98\n",
    "epsilon = 0.01\n",
    "target_update = 10\n",
    "buffer_size = 10000\n",
    "minimal_size = 500\n",
    "batch_size = 64\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "env_name = 'CartPole-v0'\n",
    "env = gym.make(env_name)\n",
    "random.seed(0)\n",
    "np.random.seed(0)\n",
    "env.seed(0)\n",
    "torch.manual_seed(0)\n",
    "state_dim = env.observation_space.shape[0]\n",
    "action_dim = env.action_space.n\n",
    "replay_buffer = rl_utils.ReplayBuffer(buffer_size)\n",
    "agent = DQN(state_dim, hidden_dim, action_dim, lr, gamma, epsilon,\n",
    "            target_update, device, 'DuelingDQN')\n",
    "return_list, max_q_value_list = train_DQN(agent, env, num_episodes,\n",
    "                                          replay_buffer, minimal_size,\n",
    "                                          batch_size)"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
